import pandas as pd
import torch
import torch.nn as nn
import numpy as np
from torch_geometric.data import Data
from torch_geometric.nn import GINConv, MLP
from sklearn.preprocessing import StandardScaler


# 定义注意力融合模块，用于融合图网络、金融数据和文本数据的特征
class AttentionFusion(nn.Module):
    def __init__(self, dim):
        super().__init__()
        # 注意力机制：将三种特征拼接后通过线性层和激活函数计算权重
        self.attention = nn.Sequential(
            nn.Linear(dim * 3, dim),
            nn.Tanh(),
            nn.Linear(dim, 3),
            nn.Softmax(dim=1)
        )

    def forward(self, gin_out, fin_out, text_out):
        # 将三种特征拼接
        combined = torch.cat([gin_out, fin_out, text_out], dim=1)
        # 计算注意力权重
        weights = self.attention(combined)
        # 对特征进行加权融合
        weighted = torch.stack([gin_out, fin_out, text_out], dim=2)
        return torch.sum(weighted * weights.unsqueeze(1), dim=2)


# 定义门控机制，用于增强特征表达
class Gate(nn.Module):
    def __init__(self, dim):
        super().__init__()
        # 门控网络：通过线性层和Sigmoid函数生成门控权重
        self.gate = nn.Sequential(nn.Linear(dim, dim), nn.Sigmoid())

    def forward(self, x):
        # 应用门控机制，增强特征
        return x * self.gate(x)


# 定义改进的三通道融合GIN模型，整合图网络、金融数据和文本数据
class ImprovedTripleChannelFusionGIN(nn.Module):
    def __init__(self, attr_dim, fin_dim, text_dim, hidden_dim=64):
        super().__init__()
        # 图网络部分：使用GINConv进行节点特征更新
        self.mlp1 = MLP([attr_dim, hidden_dim, hidden_dim], norm='batch_norm', dropout=0.2)
        self.gin1 = GINConv(self.mlp1)
        self.mlp2 = MLP([hidden_dim, hidden_dim, hidden_dim], norm='batch_norm', dropout=0.2)
        self.gin2 = GINConv(self.mlp2)
        self.mlp3 = MLP([hidden_dim, hidden_dim, 32], norm='batch_norm', dropout=0.2)
        self.gin3 = GINConv(self.mlp3)
        self.gin_bn1 = nn.BatchNorm1d(hidden_dim)
        self.gin_bn2 = nn.BatchNorm1d(hidden_dim)
        self.gin_dropout = nn.Dropout(0.2)

        # 金融数据处理分支
        self.fin_mlp = nn.Sequential(
            nn.Linear(fin_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.LeakyReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, 32)
        )
        self.fin_gate = Gate(32)

        # 文本数据处理分支
        self.text_mlp = nn.Sequential(
            nn.Linear(text_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.LeakyReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, 32)
        )
        self.text_gate = Gate(32)

        # 注意力融合模块
        self.fusion = AttentionFusion(32)

        # 分类器：将融合特征映射到最终输出
        self.classifier = nn.Sequential(
            nn.BatchNorm1d(32),
            nn.LeakyReLU(),
            nn.Dropout(0.2),
            nn.Linear(32, 16),
            nn.LeakyReLU(),
            nn.Linear(16, 1)
        )

    def forward(self, attr_x, edge_index, fin_x, text_x):
        # 图网络处理：通过三层GINConv更新节点特征
        g = torch.relu(self.gin_bn1(self.gin1(attr_x, edge_index)))
        g_res = self.gin_dropout(g)
        g = torch.relu(self.gin_bn2(self.gin2(g_res, edge_index))) + g_res
        g = self.gin3(g, edge_index)

        # 金融数据处理：通过MLP和门控机制
        f = self.fin_gate(self.fin_mlp(fin_x))
        # 文本数据处理：通过MLP和门控机制
        t = self.text_gate(self.text_mlp(text_x))

        # 融合三种特征
        fused = self.fusion(g, f, t)
        # 分类预测
        return self.classifier(fused)


# 加载测试数据
def load_test_data():
    # 读取测试数据的图关系、属性、金融数据和TF-IDF特征
    edge_df = pd.read_excel("net_test_data_relationships.xlsx")
    attr_df = pd.read_excel("net_test_data_attribute.xlsx")
    fin_df = pd.read_excel("fin_test_data.xlsx")
    tfidf_df = pd.read_csv("tfidf_test_features.csv", index_col=0)

    # 创建公司ID到索引的映射
    company_ids = sorted(set(edge_df['company_id']) | set(edge_df['related_company_id']))
    cid_map = {cid: i for i, cid in enumerate(company_ids)}

    # 构建边索引（图的邻接关系）
    edge_index = torch.tensor([
        [cid_map[c] for c in edge_df['company_id']],
        [cid_map[c] for c in edge_df['related_company_id']]
    ], dtype=torch.long)

    # 处理属性数据：过滤、排序并转换为张量
    attr_df = attr_df[attr_df['company_id'].isin(cid_map)]
    attr_df['node_idx'] = attr_df['company_id'].map(cid_map).astype(int)
    attr_df = attr_df.dropna().sort_values('node_idx')
    x_attr = torch.tensor(attr_df.drop(columns=['company_id', 'node_idx']).values, dtype=torch.float)

    # 处理金融数据：填充缺失值并标准化
    fin_df = fin_df[fin_df['company_id'].isin(cid_map)]
    fin_df['node_idx'] = fin_df['company_id'].map(cid_map).astype(int)
    fin_df = fin_df.dropna()
    fin_dim = fin_df.drop(columns=['company_id', 'node_idx']).shape[1]
    x_fin_all = np.zeros((len(cid_map), fin_dim))
    for _, row in fin_df.iterrows():
        idx = int(row['node_idx'])
        vec = row.drop(['company_id', 'node_idx']).to_numpy(dtype=np.float32)
        x_fin_all[idx] = vec
    x_fin_all = StandardScaler().fit_transform(x_fin_all)
    x_fin_tensor = torch.tensor(x_fin_all, dtype=torch.float)

    # 处理文本数据：填充TF-IDF特征
    tfidf_dim = tfidf_df.shape[1]
    x_text_all = np.zeros((len(cid_map), tfidf_dim))
    for cid, vec in tfidf_df.iterrows():
        if cid in cid_map:
            idx = cid_map[cid]
            x_text_all[idx] = vec.values
    x_text_tensor = torch.tensor(x_text_all, dtype=torch.float)

    # 生成公司ID列表
    company_id_list = [cid for cid, _ in sorted(cid_map.items(), key=lambda x: x[1])]
    # 返回图数据、金融数据张量、文本数据张量、公司ID列表和金融数据中的公司ID
    return Data(x=x_attr, edge_index=edge_index), x_fin_tensor, x_text_tensor, company_id_list, fin_df["company_id"].tolist()


# 测试模型并保存预测结果
def test_and_save():
    # 加载测试数据
    data, x_fin_tensor, x_text_tensor, all_company_ids, fin_ids = load_test_data()
    # 设置计算设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # 将数据移动到指定设备
    data = data.to(device)
    x_fin_tensor = x_fin_tensor.to(device)
    x_text_tensor = x_text_tensor.to(device)

    # 初始化模型并加载预训练权重
    model = ImprovedTripleChannelFusionGIN(
        attr_dim=data.x.shape[1],
        fin_dim=x_fin_tensor.shape[1],
        text_dim=x_text_tensor.shape[1]
    ).to(device)
    model.load_state_dict(torch.load("model_triplechannel_best.pth", map_location=device))
    model.eval()

    # 进行预测
    with torch.no_grad():
        logits = model(data.x, data.edge_index, x_fin_tensor, x_text_tensor)
        probs = torch.sigmoid(logits).squeeze().cpu().numpy()

    # 保存所有公司的预测结果
    all_preds = pd.DataFrame({
        "Company_id": all_company_ids,
        "Probability": probs
    })

    # 筛选出金融数据中存在的公司ID的预测结果
    final_results = all_preds[all_preds["Company_id"].isin(fin_ids)].reset_index(drop=True)
    # 保存最终结果到CSV文件
    final_results.to_csv("results_2022215030_朱建宇.csv", index=False)
    print("已生成三通道GIN模型测试结果：results_2022215030_朱建宇.csv")


if __name__ == "__main__":
    # 执行测试并保存结果
    test_and_save()